import pandas as pd


# 查询的列    => select colum_name1, ..., colum_namen
# 基础的选择  =>  where
def big_countries(world: pd.DataFrame) -> pd.DataFrame:
    return world[(world['area'] >= 3000000) | (world['population'] >= 25000000)][['name', 'population', 'area']]


def find_products(products: pd.DataFrame) -> pd.DataFrame:
    return products[(products['low_fats'] == 'Y') & (products['recyclable'] == 'Y')][['product_id']]


# 查询的列别名    => select colum_name1 as newname_1, ..., colum_namen as newname_n
def find_customers(customers: pd.DataFrame, orders: pd.DataFrame) -> pd.DataFrame:
    ordered_people = orders['customerId'].unique()
    return customers[~customers['id'].isin(ordered_people)].rename(columns={'name': 'Customers'})[['Customers']]


# unique
# sort
def article_views(views: pd.DataFrame) -> pd.DataFrame:
    res_indices = views['author_id'] == views['viewer_id']
    res_narr = views[res_indices].sort_values(by=['author_id'])['author_id'].unique()
    return pd.DataFrame(res_narr, columns=['id'])


# 字符相关函数
def invalid_tweets(tweets: pd.DataFrame) -> pd.DataFrame:
    invalid_indices = tweets['content'].str.len() > 15
    return tweets[invalid_indices][['tweet_id']]


def calculate_special_bonus(employees: pd.DataFrame) -> pd.DataFrame:
    has_bonus_indices = (employees['employee_id'] & 1) & (~employees['name'].str.startswith('M'))
    employees['bonus'] = employees['salary']
    employees.sort_values(by='employee_id', inplace=True)
    employees.loc[~has_bonus_indices, 'bonus'] = 0  # 这是在原视图操作
    # employees[~has_bonus_indices]['bonus'] = 0        # 这是在新视图操作， employees[~has_bonus_indices]新视图
    return employees[['employee_id', 'bonus']]


def fix_names(users: pd.DataFrame) -> pd.DataFrame:
    users['name'] = users['name'].str.capitalize()
    return users.sort_values(by=['user_id'])


# 正则
def valid_emails(users: pd.DataFrame) -> pd.DataFrame:
    return users[users["mail"].str.match(r"^[a-zA-Z][a-zA-Z0-9_.-]*\@leetcode\.com$")]


def find_patients(patients: pd.DataFrame) -> pd.DataFrame:
    condition = (patients['conditions'].str.startswith('DIAB1') | patients['conditions'].str.contains(' DIAB1'))
    return patients[condition]


# 数据操作
def nth_highest_salary(employee: pd.DataFrame, N: int) -> pd.DataFrame:
    salaries = employee['salary'].sort_values(ascending=False).unique()
    res = None if len(salaries) < N or N < 1 else salaries[N - 1]
    return pd.DataFrame([res], columns=[f'getNthHighestSalary({N})'])


def second_highest_salary(employee: pd.DataFrame) -> pd.DataFrame:
    salaries = employee['salary'].sort_values(ascending=False).unique()
    res = None if len(salaries) >= 2 else salaries[1]
    return pd.DataFrame([res], columns=[f'SecondHighestSalary'])


def department_highest_salary(employee: pd.DataFrame, department: pd.DataFrame) -> pd.DataFrame:
    df = employee.merge(department, left_on='departmentId', right_on='id', how='inner')
    df.rename(columns={'name_x': 'Employee', 'name_y': 'Department', 'salary': 'Salary'}, inplace=True)
    max_salary = df.groupby(by='Department')['Salary'].transform('max')
    df = df[df['Salary'] == max_salary]
    return df[['Employee', 'Department', 'Salary']]


def order_scores(scores: pd.DataFrame) -> pd.DataFrame:
    scores['rank'] = scores['score'].rank(method='dense', ascending=False)
    scores.sort_values(by=['rank'], ascending=True, inplace=True)
    return scores[['score', 'rank']]


def delete_duplicate_emails(person: pd.DataFrame) -> None:
    person.sort_values(by=['id'], inplace=True)
    person.drop_duplicates(subset=['email'], keep='first', inplace=True)


def rearrange_products_table(products: pd.DataFrame) -> pd.DataFrame:
    df = products.melt(id_vars='product_id', var_name='store', value_name='price')
    return df.dropna()


# 分组
def count_salary_categories(accounts: pd.DataFrame) -> pd.DataFrame:
    a = (accounts['income'] < 20000).sum()
    b = (accounts['income'] > 50000).sum()
    c = len(accounts) - a - b
    return pd.DataFrame({'category': ['Low Salary', 'Average Salary', 'High Salary'], 'accounts_count': [a, c, b]})
    # 为0的还要处理，差不多对了
    # accounts['category'] = 'Average Salary'
    # accounts.loc[accounts['income'] < 20000, 'category'] = 'Low Salary'
    # accounts.loc[accounts['income'] > 50000, 'category'] = 'High Salary'
    # df = accounts.groupby('category').size().reset_index(name = 'accounts_count')
    # return df


def total_time(employees: pd.DataFrame) -> pd.DataFrame:
    employees['total_time'] = employees['out_time'] - employees['in_time']
    df = employees.groupby(by=['emp_id', 'event_day'])["total_time"].sum().reset_index()
    return df.rename(columns={'event_day': 'day'})


def game_analysis(activity: pd.DataFrame) -> pd.DataFrame:
    return activity.groupby(by=['player_id'])['event_date'].min().reset_index(name='first_login')


def count_unique_subjects(teacher: pd.DataFrame) -> pd.DataFrame:
    return teacher.groupby(by=['teacher_id'])['subject_id'].nunique().reset_index(name='cnt')


def find_classes(courses: pd.DataFrame) -> pd.DataFrame:
    df = courses.groupby(by=['class'])['student'].nunique().reset_index()
    return df[(df['student'] > 4)][['class']]


def largest_orders(orders: pd.DataFrame) -> pd.DataFrame:
    df = orders.groupby('customer_number')['order_number'].size().reset_index(name='count')
    return df.sort_values(by=['count'], ascending=False).iloc[0:1][['customer_number']]


def categorize_products(activities: pd.DataFrame) -> pd.DataFrame:
    groupby = activities.groupby('sell_date')
    res = groupby.agg(num_sold=('product', 'nunique'),
                      products=('product', lambda x: ','.join(sorted(set(x))))).reset_index()
    return res


def daily_leads_and_partners(daily_sales: pd.DataFrame) -> pd.DataFrame:
    groupby = daily_sales.groupby(['date_id', 'make_name'])
    stats = groupby.agg(
        unique_leads=('lead_id', 'nunique'),
        unique_partners=('partner_id', 'nunique')
    ).reset_index()
    return stats


def actors_and_directors(actor_director: pd.DataFrame) -> pd.DataFrame:
    df = actor_director.groupby(['actor_id', 'director_id']).size().reset_index(name='count')
    return df[df['count'] > 2][['actor_id', 'director_id']]


# 表连接
def replace_employee_id(employees: pd.DataFrame, employee_uni: pd.DataFrame) -> pd.DataFrame:
    return employees.merge(employee_uni, left_on='id', right_on='id', how='left')[['unique_id', 'name']]


def students_and_examinations(students: pd.DataFrame, subjects: pd.DataFrame,
                              examinations: pd.DataFrame) -> pd.DataFrame:
    stu_subject = students.merge(subjects, how='cross')
    stu_subject_count = examinations.groupby(['student_id', 'subject_name']).size().reset_index(name='attended_exams')
    res = stu_subject.merge(stu_subject_count, on=['student_id', 'subject_name'], how='left')
    res['attended_exams'].fillna(0, inplace=True)
    return res.sort_values(by=['student_id', 'subject_name'])


def find_managers(employee: pd.DataFrame) -> pd.DataFrame:
    df = employee.groupby(by=['managerId']).size().reset_index(name='count')
    manager_ids = df[df['count'] > 4]['managerId']
    return employee[employee['id'].isin(manager_ids)][['name']]


def sales_person_f(sales_person: pd.DataFrame, company: pd.DataFrame, orders: pd.DataFrame) -> pd.DataFrame:
    try:
        red_company_id = company[company['name'] == 'RED']['com_id'].iloc[0]
        invalid_sales_ids = orders[orders['com_id'] == red_company_id]['sales_id']
        valid_indices = ~sales_person['sales_id'].isin(invalid_sales_ids)
        return sales_person[valid_indices][['name']]
    except:
        return sales_person[['name']]


def project_employees_i(project: pd.DataFrame, employee: pd.DataFrame) -> pd.DataFrame:
    project = project.merge(employee, left_on='employee_id', right_on='employee_id', how='left')
    return project.groupby('project_id')['experience_years'].mean().reset_index(name = 'average_years').round(2)


def rising_temperature(weather: pd.DataFrame) -> pd.DataFrame:
    # 将 recordDate 列转换为日期时间类型
    weather['recordDate'] = pd.to_datetime(weather['recordDate'])
    # 计算前一天的日期
    weather['prev_date'] = weather['recordDate'] - pd.Timedelta(days=1)
    # 将原表和包含前一天日期的表进行连接
    merged = pd.merge(weather, weather, left_on='prev_date', right_on='recordDate', suffixes=('', '_prev'))
    # 筛选出温度比前一天高的记录
    result = merged[merged['temperature'] > merged['temperature_prev']][['id']]
    return result

# data = [[1, '2015-01-01', 10], [2, '2015-01-02', 25], [3, '2015-01-03', 20], [4, '2015-01-04', 30]]
# weather = pd.DataFrame(data, columns=['id', 'recordDate', 'temperature']).astype({'id':'Int64', 'recordDate':'datetime64[ns]', 'temperature':'Int64'})
# print(rising_temperature(weather))


def trips_and_users(trips: pd.DataFrame, users: pd.DataFrame) -> pd.DataFrame:
    trips = trips.merge(users, left_on='client_id', right_on='users_id', how='left')\
                .merge(users, left_on='driver_id', right_on='users_id', how='left')
    trips = trips[(trips['banned_y'] == 'No') & (trips['banned_x'] == 'No')]
    trips = trips[(trips['request_at'] >= '2013-10-01') & (trips['request_at'] <= '2013-10-03')]
    days = []
    cancel_rate = []
    for day in ['2013-10-01', '2013-10-02', '2013-10-03']:
        day_trip = trips[trips['request_at'] == day]
        n_canceled = len(day_trip[day_trip['status'] != 'completed'])
        n_total = len(day_trip)
        if n_total:
            days.append(day)
            cancel_rate.append( round(n_canceled / n_total, 2) )
    return pd.DataFrame({'Day': days, 'Cancellation Rate': cancel_rate})

# data = [['1', '1', '10', '1', 'completed', '2013-10-01'], ['2', '2', '11', '1', 'cancelled_by_driver', '2013-10-01'], ['3', '3', '12', '6', 'completed', '2013-10-01'], ['4', '4', '13', '6', 'cancelled_by_client', '2013-10-01'], ['5', '1', '10', '1', 'completed', '2013-10-02'], ['6', '2', '11', '6', 'completed', '2013-10-02'], ['7', '3', '12', '6', 'completed', '2013-10-02'], ['8', '2', '12', '12', 'completed', '2013-10-03'], ['9', '3', '10', '12', 'completed', '2013-10-03'], ['10', '4', '13', '12', 'cancelled_by_driver', '2013-10-03']]
# trips = pd.DataFrame(data, columns=['id', 'client_id', 'driver_id', 'city_id', 'status', 'request_at']).astype({'id':'Int64', 'client_id':'Int64', 'driver_id':'Int64', 'city_id':'Int64', 'status':'object', 'request_at':'object'})
# data = [['1', 'No', 'client'], ['2', 'Yes', 'client'], ['3', 'No', 'client'], ['4', 'No', 'client'], ['10', 'No', 'driver'], ['11', 'No', 'driver'], ['12', 'No', 'driver'], ['13', 'No', 'driver']]
# users = pd.DataFrame(data, columns=['users_id', 'banned', 'role']).astype({'users_id':'Int64', 'banned':'object', 'role':'object'})
# print(trips_and_users(trips, users))


def count_occurrences(files: pd.DataFrame) -> pd.DataFrame:
    bull_cnt = files['content'].str.contains(' bull ', regex=False).sum()
    bear_cnt = files['content'].str.contains(' bear ', regex=False).sum()
    return pd.DataFrame({'word': ['bull', 'bear'], 'count': [bull_cnt, bear_cnt],})


def count_rich_customers(store: pd.DataFrame) -> pd.DataFrame:
    df = store.groupby('customer_id')['amount'].max().reset_index(name='max_amount')
    cnt = (df['max_amount'] > 500).sum()
    return pd.DataFrame({'rich_count': [cnt]})

def food_delivery(delivery: pd.DataFrame) -> pd.DataFrame:
    all_cnt = len(delivery)
    immediate_cnt = (delivery['order_date'] == delivery['customer_pref_delivery_date']).sum()
    return pd.DataFrame({'immediate_percentage': [round(immediate_cnt / all_cnt * 100, 2)]})

def ads_performance(ads: pd.DataFrame) -> pd.DataFrame:
    groupby = ads.groupby('ad_id')
    stats = groupby.agg(
        n_clicked=('action', lambda x: (x=='Clicked').sum()),
        n_viewed=('action', lambda x: (x=='Viewed').sum())
    ).reset_index()
    stats['ctr'] = .0
    indices = (stats['n_clicked'] + stats['n_viewed']) != 0
    stats.loc[indices, 'ctr'] = (stats['n_clicked'] / (stats['n_clicked'] + stats['n_viewed']) * 100).round(2)
    stats.sort_values(by = ['ctr','ad_id'], ascending=[False, True], inplace=True)
    return stats[:][['ad_id', 'ctr']]

# data = [[1, 1, 'Clicked'], [2, 2, 'Clicked'], [3, 3, 'Viewed'], [5, 5, 'Ignored'], [1, 7, 'Ignored'], [2, 7, 'Viewed'], [3, 5, 'Clicked'], [1, 4, 'Viewed'], [2, 11, 'Viewed'], [1, 2, 'Clicked']]
# ads = pd.DataFrame(data, columns=['ad_id', 'user_id', 'action']).astype({'ad_id':'Int64', 'user_id':'Int64', 'action':'object'})
# print(ads_performance(ads))


def accepted_candidates(candidates: pd.DataFrame, rounds: pd.DataFrame) -> pd.DataFrame:
    df = rounds.groupby('interview_id')['score'].sum()
    df = candidates.merge(df, left_on='interview_id', right_on='interview_id', how='left')
    indices = (df['score'] > 15) & (df['years_of_exp'] >= 2)
    return df[indices][['candidate_id']]

data = [[11, 'Atticus', 1, 101], [9, 'Ruben', 6, 104], [6, 'Aliza', 10, 109], [8, 'Alfredo', 0, 107]]
candidates = pd.DataFrame(data, columns=['candidate_id', 'name', 'years_of_exp', 'interview_id']).astype({'candidate_id':'Int64', 'name':'object', 'years_of_exp':'Int64', 'interview_id':'Int64'})
data = [[109, 3, 4], [101, 2, 8], [109, 4, 1], [107, 1, 3], [104, 3, 6], [109, 1, 4], [104, 4, 7], [104, 1, 2], [109, 2, 1], [104, 2, 7], [107, 2, 3], [101, 1, 8]]
rounds = pd.DataFrame(data, columns=['interview_id', 'round_id', 'score']).astype({'interview_id':'Int64', 'round_id':'Int64', 'score':'Int64'})
print(accepted_candidates(candidates, rounds))


